package com.hyz.yuaiagent.rag.document;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.transformer.KeywordMetadataEnricher;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * MyKeywordEnricher组件用于 enrich 文档列表中的关键词
 */
@Component
public class MyKeywordEnricher {

    // 使用的聊天模型，用于关键词提取
    private final ChatModel dashscopeChatModel;

    /**
     * 构造函数注入ChatModel
     *
     * @param dashscopeChatModel 聊天模型，用于关键词提取
     */
    @Autowired
    public MyKeywordEnricher(ChatModel dashscopeChatModel) {
        this.dashscopeChatModel = dashscopeChatModel;
    }

    /**
     * 对文档列表进行关键词 enrich 操作
     *
     * @param documents 文档列表，需要进行关键词 enrich
     * @return 返回 enrich 过的文档列表
     */
    public List<Document> KeywordEnricher(List<Document> documents) {
        // 创建KeywordMetadataEnricher实例，指定聊天模型和关键词数量
        KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(dashscopeChatModel, 5);

        // 应用KeywordMetadataEnricher对文档列表进行关键词 enrich
        return keywordMetadataEnricher.apply(documents);
    }
}
